package org.example.rmi;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManagerFactory;
import javax.rmi.ssl.SslRMIClientSocketFactory;
import java.io.IOException;
import java.net.Socket;
import java.rmi.RemoteException;
import java.security.KeyStore;
import java.security.SecureRandom;

public class MyRMIClientSocketFactory extends SslRMIClientSocketFactory {

    private static final long serialVersionUID = 1L;

    @Override
    public Socket createSocket(String host, int port) throws IOException {
        try {
            SSLContext sslContext = getSSLContext();
            SSLSocketFactory socketFactory = sslContext.getSocketFactory();
            return socketFactory.createSocket(host, port);
        } catch (Exception e) {
//            e.printStackTrace();
            throw new RemoteException();
        }
    }

    /**
     * keytool -genkeypair -alias server -keyalg RSA -keystore server-keystore.jks
     * keytool -export -alias server -keystore server-keystore.jks -file server-cert.crt
     * keytool -import -alias server -file server-cert.crt -keystore client-truststore.jks
     */
    public static SSLContext getSSLContext() throws Exception {
        String keyStorePath = "server-keystore.jks";
        String trustStorePath = "client-truststore.jks";
        String storepwd = "password";

        KeyStore keyKeyStore = KeyStore.getInstance("JKS");
        keyKeyStore.load(MyRMIClientSocketFactory.class.getClassLoader().getResourceAsStream(keyStorePath), storepwd.toCharArray());
        KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
        keyManagerFactory.init(keyKeyStore, storepwd.toCharArray());

        KeyStore trustStore = KeyStore.getInstance("JKS");
        trustStore.load(MyRMIClientSocketFactory.class.getClassLoader().getResourceAsStream(trustStorePath), storepwd.toCharArray());
        TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
        trustManagerFactory.init(trustStore);

        SSLContext sslContext = SSLContext.getInstance("TLSv1.2");
        sslContext.init(keyManagerFactory.getKeyManagers(), trustManagerFactory.getTrustManagers(), new SecureRandom());
        return sslContext;
    }
}
